"""
执行计划生成器
将AST转换为执行计划，输出CreateTable、Insert、SeqScan、Filter、Project算子
"""

import json

class ExecutionPlan:
    """执行计划类"""
    
    def __init__(self, plan_type, **kwargs):
        self.plan_type = plan_type
        self.attributes = kwargs
        self.children = []
    
    def add_child(self, child_plan):
        """添加子计划"""
        self.children.append(child_plan)
    
    def to_dict(self):
        """转换为字典格式"""
        result = {
            'type': self.plan_type,
            'attributes': self.attributes
        }
        if self.children:
            result['children'] = [child.to_dict() for child in self.children]
        return result
    
    def to_json(self):
        """转换为JSON格式"""
        return json.dumps(self.to_dict(), indent=2, ensure_ascii=False)
    
    def to_s_expression(self):
        """转换为S表达式格式"""
        if not self.children:
            attrs = ' '.join([f"{k}={v}" for k, v in self.attributes.items()])
            return f"({self.plan_type} {attrs})"
        else:
            children_str = ' '.join([child.to_s_expression() for child in self.children])
            attrs = ' '.join([f"{k}={v}" for k, v in self.attributes.items()])
            return f"({self.plan_type} {attrs} {children_str})"

class ExecutionPlanGenerator:
    """执行计划生成器"""
    
    def __init__(self):
        """初始化执行计划生成器"""
        pass
    
    def generate_plan(self, ast):
        """根据AST生成执行计划"""
        if not ast:
            return None
        
        if ast['type'] == 'CREATE_TABLE':
            return self.generate_create_table_plan(ast)
        elif ast['type'] == 'INSERT':
            return self.generate_insert_plan(ast)
        elif ast['type'] == 'SELECT':
            return self.generate_select_plan(ast)
        elif ast['type'] == 'UPDATE':
            return self.generate_update_plan(ast)
        elif ast['type'] == 'DELETE':
            return self.generate_delete_plan(ast)
        else:
            return ExecutionPlan("ERROR", message=f"不支持的语句类型: {ast['type']}")
    
    def generate_create_table_plan(self, ast):
        """生成CREATE TABLE执行计划"""
        # 从AST构建schema
        schema = {
            'columns': {},
            'primary_key': None
        }
        
        for col_def in ast['columns']:
            col_name = col_def['name']
            schema['columns'][col_name] = col_def
            if col_def.get('primary_key'):
                schema['primary_key'] = col_name
        
        return ExecutionPlan(
            "CreateTable",
            table_name=ast['table_name'],
            schema=schema
        )
    
    def generate_insert_plan(self, ast):
        """生成INSERT执行计划"""
        return ExecutionPlan(
            "Insert",
            table_name=ast['table_name'],
            values=ast.get('values', []),
            columns=ast.get('columns', [])
        )
    
    def generate_select_plan(self, ast):
        """生成SELECT执行计划"""
        # 创建Project算子
        project_plan = ExecutionPlan(
            "Project",
            columns=ast['columns']
        )
        
        # 如果有WHERE条件，创建Filter算子
        if ast['where']:
            filter_plan = ExecutionPlan(
                "Filter",
                condition=ast['where']
            )
            filter_plan.add_child(project_plan)
            
            # 创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            filter_plan.add_child(seqscan_plan)
            
            return filter_plan
        else:
            # 没有WHERE条件，直接创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            project_plan.add_child(seqscan_plan)
            
            return project_plan
    
    def generate_update_plan(self, ast):
        """生成UPDATE执行计划"""
        # 创建Update算子
        update_plan = ExecutionPlan(
            "Update",
            table_name=ast['table_name'],
            assignments=ast['assignments']
        )
        
        # 如果有WHERE条件，创建Filter算子
        if ast['where']:
            filter_plan = ExecutionPlan(
                "Filter",
                condition=ast['where']
            )
            filter_plan.add_child(update_plan)
            
            # 创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            filter_plan.add_child(seqscan_plan)
            
            return filter_plan
        else:
            # 没有WHERE条件，直接创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            update_plan.add_child(seqscan_plan)
            
            return update_plan
    
    def generate_delete_plan(self, ast):
        """生成DELETE执行计划"""
        # 创建Delete算子
        delete_plan = ExecutionPlan(
            "Delete",
            table_name=ast['table_name']
        )
        
        # 如果有WHERE条件，创建Filter算子
        if ast['where']:
            filter_plan = ExecutionPlan(
                "Filter",
                condition=ast['where']
            )
            filter_plan.add_child(delete_plan)
            
            # 创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            filter_plan.add_child(seqscan_plan)
            
            return filter_plan
        else:
            # 没有WHERE条件，直接创建SeqScan算子
            seqscan_plan = ExecutionPlan(
                "SeqScan",
                table_name=ast['table_name']
            )
            delete_plan.add_child(seqscan_plan)
            
            return delete_plan
    
    def print_plan_tree(self, plan, indent=0):
        """打印执行计划树"""
        if not plan:
            return
        
        print("  " * indent + f"{plan.plan_type}")
        for key, value in plan.attributes.items():
            print("  " * (indent + 1) + f"{key}: {value}")
        
        for child in plan.children:
            self.print_plan_tree(child, indent + 1)

# 测试函数
if __name__ == "__main__":
    generator = ExecutionPlanGenerator()
    
    # 测试用例
    test_asts = [
        {
            'type': 'CREATE_TABLE',
            'table_name': 'student',
            'schema': {
                'columns': {
                    'id': {'type': {'type': 'int'}},
                    'name': {'type': {'type': 'varchar', 'length': 50}},
                    'age': {'type': {'type': 'int'}}
                },
                'primary_key': 'id'
            }
        },
        {
            'type': 'INSERT',
            'table_name': 'student',
            'values': [1, 'Alice', 20],
            'columns': ['id', 'name', 'age']
        },
        {
            'type': 'SELECT',
            'table_name': 'student',
            'columns': '*',
            'where': {
                'type': 'GT',
                'left': 'age',
                'right': 18
            }
        },
        {
            'type': 'DELETE',
            'table_name': 'student',
            'where': {
                'type': 'EQ',
                'left': 'id',
                'right': 1
            }
        }
    ]
    
    for i, ast in enumerate(test_asts):
        print(f"\n测试用例 {i+1}: {ast['type']}")
        print("-" * 50)
        
        plan = generator.generate_plan(ast)
        if plan:
            print("执行计划树:")
            generator.print_plan_tree(plan)
            
            print("\nJSON格式:")
            print(plan.to_json())
            
            print("\nS表达式格式:")
            print(plan.to_s_expression())
